# TODO: 导入必要的库和模块
import numpy as np  # 用于数值计算和数组操作
import pickle  # 用于模型的序列化和存储
from PIL import Image  # 用于图像处理
import gradio as gr  # 用于创建用户交互界面
import io  # 用于处理字节流

# TODO: 加载保存的KNN模型，这样我们可以使用预训练的模型进行预测
with open('best_knn_model.pkl', 'rb') as f:  # 以二进制读取模式打开保存的KNN模型文件
    knn_model = pickle.load(f)  # 加载模型

# TODO: 定义预测函数，这个函数将用于Gradio接口进行预测
def predict_digit(image):
    print("Received image data:", image)  # 调试输出：打印收到的图像数据

    # 检查图像数据格式
    if isinstance(image, dict):
        if "composite" in image:  # 如果传入的是字典并且包含"composite"键
            image = image["composite"]  # 提取图像
        else:
            print("No composite image found.")  # 调试输出：没有找到复合图像
            return None

    image = np.array(image)  # 将图像转换为NumPy数组

    # 检查图像数据类型并进行转换
    if image.dtype != np.uint8:  # 如果类型不是uint8
        image = (image * 255).astype(np.uint8)  # 将图像数据缩放到0-255并转换类型

    try:
        # 将图像转换为灰度，调整为8x8大小
        image = Image.fromarray(image).convert('L').resize((8, 8))  # 调整为8x8
    except Exception as e:
        print("Error converting image:", e)  # 输出转换错误信息
        return None

    # 将图像数组变形为（1, -1），因为8x8有64个像素
    img_array = np.array(image).reshape(1, -1)  # 变形为（1, 64）
    img_array = (img_array / 255) * 16  # 归一化处理

    # 添加调试信息：输出处理后的图像数组形状和数值
    print("Processed image array shape:", img_array.shape)
    print("Processed image array values:", img_array)

    # 使用训练好的KNN模型进行预测
    prediction = knn_model.predict(img_array)  # 进行预测
    print("Prediction:", prediction)  # 输出预测结果

    return int(prediction[0])  # 返回预测的数字

# TODO: 创建Gradio接口，这个接口将用于用户输入和显示预测结果
with gr.Blocks() as demo:  # 创建Gradio的块界面
    gr.Markdown("## 手写数字识别器")  # 添加Markdown标题

    with gr.Row():  # 创建一行
        drawing_input = gr.Sketchpad(label="在这里绘制手写数字")  # 绘图输入区
        prediction_output = gr.Label(label="预测的数字:")  # 显示预测结果的标签

    submit_btn = gr.Button("提交")  # 提交按钮
    clear_btn = gr.Button("清除")  # 清除按钮

    # 处理提交的图像
    def process_image(img):
        if img is None or len(img) == 0:  # 检查是否有绘制内容
            return "请绘制一个数字"  # 提示用户输入
        return predict_digit(img)  # 调用预测函数并返回结果

    # 清除画布
    def clear_canvas():
        return None  # 返回None，清空画布

    # 设置按钮的点击事件
    submit_btn.click(fn=process_image, inputs=drawing_input, outputs=prediction_output)  # 提交按钮点击事件
    clear_btn.click(fn=clear_canvas, inputs=[], outputs=drawing_input)  # 清除按钮点击事件

# TODO: 启动Gradio接口，用户可以通过这个接口进行交互
demo.launch()  # 启动Gradio界面，用户可以进行数字识别交互。如果需要公共链接，可以设置share=True